#include "pm.h"

static  int patch_w  = 7;
static  int pm_iters = 5;
static  int rs_max   = INT_MAX;


int dist( pnm a, pnm b, int ai, int aj, int bi, int bj, int cutoff ){

  int result = 0, i = 0, j = 0;
  for( i = 0; i < patch_w; i++ ){
    for ( j = 0; j < patch_w; j++ ){
      // récupérer les valeurs des canaux de l'image a
      int ra = pnm_get_component( a, ai + i, aj + j, 0);
      int ga = pnm_get_component( a, ai + i, aj + j, 1);
      int ba = pnm_get_component( a, ai + i, aj + j, 2);
      // récupérer les valeurs des canaux de l'image b
      int rb = pnm_get_component( b, bi + i, bj + j, 0);
      int gb = pnm_get_component( b, bi + i, bj + j, 1);
      int bb = pnm_get_component( b, bi + i, bj + j, 2);
      
      int dr = ra - rb, dg = ga - gb, db = ba - bb;

      result += dr * dr + dg * dg + db * db;
    }
    if ( result >= cutoff ) 
      return cutoff;
  }

  return result;
}

void improve_guess( pnm a, pnm b, int ai, int aj, int *ibest, int *jbest, int *dbest, int bi, int bj) {
  int d = dist(a, b, ai, aj, bi, bj, *dbest);
  if (d < *dbest) {
    *dbest = d;
    *ibest = bi;
    *jbest = bj;
  }
}

/* Match image a to image b, returning the nearest neighbor field mapping a => b coords, stored in an RGB 24-bit image as (by<<12)|bx. */
void patchmatch( pnm a, pnm b, int * ann, int * annd) {
  /* Initialize with random nearest neighbor field (NNF). */
  // ann = new BITMAP(a->w, a->h);
  // annd = new BITMAP(a->w, a->h);
  int width = pnm_get_width(a);
  int height = pnm_get_height(a);

  int aew = width - patch_w+1, aeh = height - patch_w + 1;       /* Effective width and height (possible upper left corners of patches). */
  int bew = pnm_get_width(b) - patch_w+1, beh = pnm_get_height(b) - patch_w + 1;

  
  // memset(ann->data, 0, sizeof(int)*a->w*a->h);
  // memset(annd->data, 0, sizeof(int)*a->w*a->h);
  int ai = 0, aj = 0;
  for ( ai = 0; ai < aeh; ai++ ) {
    for ( aj = 0; aj < aew; aj++ ) {
      int bi = rand()%beh;
      int bj = rand()%bew;
      
      ann[ai * width + aj] = IJ_TO_INT(bj, bi);
      annd[ai * width + aj] = dist(a, b, ai, aj, bi, bj, INT_MAX);
    }
  }

  int iter = 0;
  for ( iter = 0; iter < pm_iters; iter++) {
    /* In each iteration, improve the NNF, by looping in scanline or reverse-scanline order. */
    int istart = 0, iend = aeh, ichange = 1;
    int jstart = 0, jend = aew, jchange = 1;
    if (iter % 2 == 1) {
      jstart = jend-1; jend = -1; jchange = -1;
      istart = iend-1; iend = -1; ichange = -1;
    }
    for ( ai = istart; ai != iend; ai += ichange ) {
      for ( aj = jstart; aj != jend; aj += jchange ) { 
        /* Current (best) guess. */
        int v = ann[ai * width + aj];
        int jbest = INT_TO_J(v), ibest = INT_TO_I(v);
        int dbest = annd[ai * width + aj];

        /* Propagation: Improve current guess by trying instead correspondences from left and above (below and right on odd iterations). */
        if ((unsigned) (aj - jchange) < (unsigned) aew) {
          int vp = ann[ai * width + aj - jchange];
          int jp = INT_TO_J(vp) + jchange, ip = INT_TO_I(vp);
          if ((unsigned) jp < (unsigned) bew) {
            improve_guess(a, b, ai, aj, &ibest, &jbest, &dbest, ip, jp);
          }
        }

        if ((unsigned) (ai - ichange) < (unsigned) aeh) {
          int vp = ann[(ai - ichange) * width + aj];
          int jp = INT_TO_J(vp), ip = INT_TO_I(vp) + ichange;
          if ((unsigned) ip < (unsigned) beh) {
            improve_guess(a, b, ai, aj, &ibest, &jbest, &dbest, ip, jp);
          }
        }

        /* Random search: Improve current guess by searching in boxes of exponentially decreasing size around the current best guess. */
        int rs_start = rs_max;
        if (rs_start > MAX( pnm_get_width(b), pnm_get_height(b))) { rs_start = MAX( pnm_get_width(b), pnm_get_height(b)); }
        int mag = rs_start;
        for ( mag = rs_start; mag >= 1; mag /= 2 ) {
          /* Sampling window */
          int jmin = MAX(jbest-mag, 0), jmax = MIN(jbest+mag+1,bew);
          int imin = MAX(ibest-mag, 0), imax = MIN(ibest+mag+1,beh);
          int jp = jmin+rand()%(jmax-jmin);
          int ip = imin+rand()%(imax-imin);
          improve_guess(a, b, ai, aj, &ibest, &jbest, &dbest, ip, jp);
        }

        ann[ai * width + aj] = IJ_TO_INT(jbest, ibest);
        annd[ai * width + aj] = dbest;
      }
    }
  }
}



int main(int argc, char *argv[] ){
  
  // int patch_w  = 7;
  // int pm_iters = 5;
  // int rs_max   = INT_MAX;

  int i, j;

  argc--;
  argv++;
  if (argc != 4) { fprintf(stderr, "pm_minimal a b ann annd\n"
                                   "Given input images a, b outputs nearest neighbor field 'ann' mapping a => b coords, and the squared L2 distance 'annd'\n"
                                   "These are stored as RGB 24-bit images, with a 24-bit int at every pixel. For the NNF we store (by<<12)|bx."); exit(1); }
  printf("Loading input images\n");


  pnm a = pnm_load(argv[0]);
  pnm b = pnm_load(argv[1]);

  int width = pnm_get_width(a);
  int height = pnm_get_height(a);


  pnm annPnm = pnm_new(width, height, PnmRawPpm);
  pnm anndPnm = pnm_new(width, height, PnmRawPpm);
  int * ann = NULL, * annd = NULL;
  
  ann = malloc( width * height * sizeof(int) );
  annd = malloc( width * height * sizeof(int) );
  
  printf("Running PatchMatch\n");
  patchmatch(a, b, ann, annd);
  printf("Saving output images\n");
  
  for( i = 0; i < height; i++ ){
    for( j = 0; j < width; j++ ){
      pnm_set_component( annPnm, i, j, 0,  ann[ i * width + j ] & 255 );
      pnm_set_component( annPnm, i, j, 1, (ann[ i * width + j ] >> 8) & 255);
      pnm_set_component( annPnm, i, j, 2, (ann[ i * width + j ] >> 16) & 255);

      pnm_set_component( anndPnm, i, j, 0,   annd[ i * width + j ] & 255);
      pnm_set_component( anndPnm, i, j, 1,  (annd[ i * width + j ] >> 8 )& 255);
      pnm_set_component( anndPnm, i, j, 2,  (annd[ i * width + j ] >> 16)& 255);
    }
  }

  pnm_save(annPnm, PnmRawPpm, argv[2]);
  pnm_save(anndPnm , PnmRawPpm, argv[3]);

  pnm_free(annPnm);
  pnm_free(anndPnm);
  pnm_free(a);
  pnm_free(b);

	return 0;
}